//
//  PredictionHelper.swift
//  MyCloset
//
//  Created by 牟好好 on 2024/10/28.
//

import UIKit
import CoreML

// 模型预测逻辑
func performStylesPrediction(index: Int, onImage image: UIImage, completion: @escaping ([String: Double]) -> Void) {
    if let pixelBuffer = convertUIImageToCVPixelBuffer(image: image) {
        let allResults: [String: Double]
        do {
            let configuration = MLModelConfiguration()
            // 根据下标调用不同模型
            switch index {
            case 0:
                // 上装
                let model = try ShangZhuangClothesStyle(configuration: configuration)
                let prediction = try model.prediction(image: pixelBuffer)
                allResults = prediction.targetProbability
            case 1:
                let model = try XiaZhuangClothesStyle(configuration: configuration)
                let prediction = try model.prediction(image: pixelBuffer)
                allResults = prediction.targetProbability
            case 2:
                let model = try WaiTaoClothesStyle(configuration: configuration)
                let prediction = try model.prediction(image: pixelBuffer)
                allResults = prediction.targetProbability
            case 3:
                let model = try XieZiClothesStyle(configuration: configuration)
                let prediction = try model.prediction(image: pixelBuffer)
                allResults = prediction.targetProbability
            case 4:
                let model = try MaoZiClothesStyle(configuration: configuration)
                let prediction = try model.prediction(image: pixelBuffer)
                allResults = prediction.targetProbability
            default:
                print("无效的模型索引")
                return
            }
            completion(allResults)
        } catch {
            print("模型预测出错: \(error.localizedDescription)")
        }
    }
}

func performTypesPrediction(onImage image: UIImage, completion: @escaping ([String: Double]) -> Void) {
    if let pixelBuffer = convertUIImageToCVPixelBuffer(image: image) {
        do {
            let configuration = MLModelConfiguration()
            let model = try ClothesTypes(configuration: configuration) // 模型
            let prediction = try model.prediction(image: pixelBuffer)
            let allResults = prediction.targetProbability
            completion(allResults)
        } catch {
            print("模型预测出错: \(error.localizedDescription)")
        }
    }
}

// Helper 函数：将 UIImage 转换为 CVPixelBuffer
func convertUIImageToCVPixelBuffer(image: UIImage) -> CVPixelBuffer? {
    let width = Int(image.size.width)
    let height = Int(image.size.height)
    
    var pixelBuffer: CVPixelBuffer?
    
    // 设置 PixelBuffer 的属性
    let attributes: [NSObject: AnyObject] = [
        kCVPixelBufferCGImageCompatibilityKey: kCFBooleanTrue,
        kCVPixelBufferCGBitmapContextCompatibilityKey: kCFBooleanTrue
    ]
    
    // 创建 PixelBuffer
    let status = CVPixelBufferCreate(kCFAllocatorDefault, width, height, kCVPixelFormatType_32ARGB, attributes as CFDictionary, &pixelBuffer)
    
    guard status == kCVReturnSuccess, let buffer = pixelBuffer else {
        print("CVPixelBuffer 创建失败")
        return nil
    }
    
    // 锁定 PixelBuffer 的基地址
    CVPixelBufferLockBaseAddress(buffer, .readOnly)
    
    // 获取 PixelBuffer 的基地址
    let pixelData = CVPixelBufferGetBaseAddress(buffer)
    
    // 创建颜色空间
    let colorSpace = CGColorSpaceCreateDeviceRGB()
    
    // 创建 CGContext，用于将 UIImage 绘制到 PixelBuffer
    let context = CGContext(
        data: pixelData,
        width: width,
        height: height,
        bitsPerComponent: 8,
        bytesPerRow: CVPixelBufferGetBytesPerRow(buffer),
        space: colorSpace,
        bitmapInfo: CGImageAlphaInfo.noneSkipFirst.rawValue
    )
    
    guard let cgImage = image.cgImage else {
        print("UIImage 转换为 CGImage 失败")
        return nil
    }
    
    // 将 UIImage 绘制到 PixelBuffer
    context?.draw(cgImage, in: CGRect(x: 0, y: 0, width: width, height: height))
    
    // 解锁 PixelBuffer
    CVPixelBufferUnlockBaseAddress(buffer, .readOnly)
    
    return buffer
}
